import os
import torch
from torch.utils.data import Dataset, DataLoader
from typing import List, Dict, Optional
import random

class IMDBDataset(Dataset):
    """PyTorch Dataset for IMDB Movie Reviews"""
    
    def __init__(self, dataset_path: str, split: str = 'test', max_samples: Optional[int] = None):
        """
        Initialize IMDB dataset
        Args:
            dataset_path: Path to IMDB dataset directory
            split: 'train' or 'test'
            max_samples: Maximum number of samples to load (None for all)
        """
        self.dataset_path = dataset_path
        self.split = split
        self.max_samples = max_samples
        self.reviews = self._load_reviews()
    
    def _load_reviews(self) -> List[Dict]:
        """Load reviews from IMDB dataset files"""
        reviews = []
        
        # Load positive reviews (label 1)
        positive_dir = os.path.join(self.dataset_path, 'aclImdb', self.split, 'pos')
        if os.path.exists(positive_dir):
            pos_reviews = self._load_reviews_from_dir(positive_dir, label=1)
            reviews.extend(pos_reviews)
        
        # Load negative reviews (label 0)
        negative_dir = os.path.join(self.dataset_path, 'aclImdb', self.split, 'neg')
        if os.path.exists(negative_dir):
            neg_reviews = self._load_reviews_from_dir(negative_dir, label=0)
            reviews.extend(neg_reviews)
        
        # Shuffle and limit samples if specified
        random.shuffle(reviews)
        if self.max_samples:
            reviews = reviews[:self.max_samples]
        
        return reviews
    
    def _load_reviews_from_dir(self, directory: str, label: int) -> List[Dict]:
        """Load reviews from a specific directory"""
        reviews = []
        try:
            for filename in os.listdir(directory):
                if filename.endswith('.txt'):
                    filepath = os.path.join(directory, filename)
                    with open(filepath, 'r', encoding='utf-8') as f:
                        text = f.read().strip()
                        reviews.append({
                            'text': text,
                            'original_label': label,
                            'target_label': 1 - label  # Flip label for attack
                        })
        except Exception as e:
            print(f"Error loading reviews from {directory}: {e}")
        
        return reviews
    
    def __len__(self):
        return len(self.reviews)
    
    def __getitem__(self, idx):
        return self.reviews[idx]

def create_imdb_dataloader(dataset_path: str, split: str = 'test', 
                          max_samples: Optional[int] = None, 
                          batch_size: int = 32, shuffle: bool = True) -> DataLoader:
    """
    Create a PyTorch DataLoader for IMDB dataset
    Args:
        dataset_path: Path to IMDB dataset directory
        split: 'train' or 'test'
        max_samples: Maximum number of samples to load
        batch_size: Batch size for the DataLoader
        shuffle: Whether to shuffle the data
    Returns:
        PyTorch DataLoader
    """
    dataset = IMDBDataset(dataset_path, split, max_samples)
    dataloader = DataLoader(
        dataset, 
        batch_size=batch_size, 
        shuffle=shuffle,
        collate_fn=lambda batch: batch  # Return raw batch for custom processing
    )
    return dataloader

def load_json_data(json_path: str) -> List[Dict]:
    """Load data from JSON file"""
    import json
    try:
        with open(json_path, 'r') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"Data file not found: {json_path}")
        return []
    except json.JSONDecodeError:
        print(f"Invalid JSON format in: {json_path}")
        return []

if __name__ == "__main__":
    # Test the dataset
    dataset = IMDBDataset("databases/imdb/imdb", split='test', max_samples=5)
    print(f"Dataset size: {len(dataset)}")
    for i, sample in enumerate(dataset):
        print(f"Sample {i}: {sample['text'][:100]}... (label: {sample['original_label']}, target: {sample['target_label']})")
